package actor

import (
	"context"
	"sync"
	"time"

	"github.com/asynkron/protoactor-go/actor"
	"github.com/sirupsen/logrus"

	"xagent/internal/knowledge"
	"xagent/internal/memory"
	proto "xagent/proto/v1"
)

// Message represents a message in the system
type Message struct {
	ID        string
	Type      string
	Content   string
	Metadata  map[string]string
	From      string
	To        string
	CreatedAt time.Time
	UpdatedAt time.Time
	State     string
	Priority  int
	Tags      []string
}

// MessagePool manages message routing
type MessagePool struct {
	messages map[string]*Message
	mutex    sync.RWMutex
}

// NewMessagePool creates a new message pool
func NewMessagePool() *MessagePool {
	return &MessagePool{
		messages: make(map[string]*Message),
	}
}

// AddMessage adds a message to the pool
func (p *MessagePool) AddMessage(msg *Message) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.messages[msg.ID] = msg
}

// GetMessage gets a message from the pool
func (p *MessagePool) GetMessage(id string) (*Message, bool) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	msg, ok := p.messages[id]
	return msg, ok
}

// EnvironmentActor represents the environment for all agents
type EnvironmentActor struct {
	*BaseActor
	members     map[string]*actor.PID
	memory      memory.MemoryManager
	knowledge   *knowledge.RAGEngine
	messagePool *MessagePool
	mutex       sync.RWMutex
	logger      *logrus.Logger
}

// NewEnvironmentActor creates a new environment actor
func NewEnvironmentActor(memory memory.MemoryManager, knowledge *knowledge.RAGEngine) *EnvironmentActor {
	return &EnvironmentActor{
		BaseActor:   NewBaseActor("environment", "Environment", "Global environment for all agents"),
		members:     make(map[string]*actor.PID),
		memory:      memory,
		knowledge:   knowledge,
		messagePool: NewMessagePool(),
		logger:      logrus.New(),
	}
}

// Receive handles incoming messages
func (e *EnvironmentActor) Receive(context actor.Context) {
	e.BaseActor.Receive(context)

	switch msg := context.Message().(type) {
	case *proto.AgentMessage:
		e.handleAgentMessage(context, msg)
	}
}

// handleAgentMessage handles messages between agents
func (e *EnvironmentActor) handleAgentMessage(ctx actor.Context, msg *proto.AgentMessage) {
	// Create environment message
	envMsg := &Message{
		ID:        msg.Id,
		Type:      msg.Type.String(),
		Content:   e.getMessageContent(msg),
		Metadata:  msg.Metadata,
		From:      msg.Sender,
		To:        msg.Receiver,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
		State:     "ACTIVE",
		Priority:  1,
		Tags:      []string{},
	}

	// Store in message pool
	e.messagePool.AddMessage(envMsg)

	// Route message
	if msg.Receiver != "" {
		if pid, ok := e.members[msg.Receiver]; ok {
			ctx.Send(pid, msg)
		}
	}
}

// Helper methods
func (e *EnvironmentActor) getMessageContent(msg *proto.AgentMessage) string {
	switch c := msg.Content.(type) {
	case *proto.AgentMessage_Message:
		return c.Message
	case *proto.AgentMessage_Error:
		return c.Error.Message
	default:
		return ""
	}
}

// Environment state management
func (e *EnvironmentActor) GetState() map[string]interface{} {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	state := make(map[string]interface{})
	state["member_count"] = len(e.members)
	state["message_count"] = len(e.messagePool.messages)
	state["last_updated"] = time.Now()
	return state
}

// Memory management
func (e *EnvironmentActor) StoreMemory(ctx context.Context, entry *memory.MemoryEntry) error {
	return e.memory.Store(ctx, entry)
}

func (e *EnvironmentActor) RetrieveMemory(ctx context.Context, id string) (*memory.MemoryEntry, error) {
	return e.memory.Get(ctx, id)
}

// Knowledge management
func (e *EnvironmentActor) Search(ctx context.Context, query string, limit int) ([]knowledge.SearchMatch, error) {
	return e.knowledge.Search(ctx, query, limit)
}

func (e *EnvironmentActor) AddDocument(ctx context.Context, doc *knowledge.Document) error {
	return e.knowledge.AddDocument(ctx, doc)
}
